#include <gtest/gtest.h>
#include "Thread/WorkThreadPool.h"
#include <chrono>
#include <atomic>
#include <vector>
#include <future>

using namespace toolkit;
using namespace std::chrono_literals;

// 测试 WorkThread 基本功能
TEST(WorkThreadTest, BasicFunctionality) {
    WorkThread thread;
    thread.start();
    
    // 初始负载应为0
    EXPECT_EQ(thread.getLoad(), 0);
    
    std::atomic<bool> task_executed{false};
    
    // 提交任务
    auto cancelable = thread.async([&task_executed]() {
        task_executed = true;
        std::this_thread::sleep_for(50ms);
    });
    
    // 等待任务执行完成
    std::this_thread::sleep_for(100ms);
    
    EXPECT_TRUE(task_executed);
    
    // 停止线程
    thread.stop();
}

// 测试 WorkThread 的同步执行功能
TEST(WorkThreadTest, SyncExecution) {
    WorkThread thread;
    thread.start();
    
    std::atomic<bool> in_work_thread{false};
    std::atomic<bool> task_executed{false};
    
    // 在工作线程中执行任务，检查是否同步执行
    thread.async([&]() {
        in_work_thread = thread.isCurrentThread();
        
        // 在工作线程中提交任务，应该同步执行
        thread.async([&task_executed]() {
            task_executed = true;
        }, true);
        
        // 任务应该已经同步执行完成
        EXPECT_TRUE(task_executed);
    });
    
    // 等待任务执行完成
    std::this_thread::sleep_for(100ms);
    
    EXPECT_TRUE(in_work_thread);
    
    thread.stop();
}

// 测试 WorkThreadPool 基本功能
TEST(WorkThreadPoolTest, BasicFunctionality) {
    WorkThreadPool pool(4); // 创建4个线程的工作线程池
    
    // 初始负载应为0
    EXPECT_EQ(pool.getLoad(), 0);
    EXPECT_EQ(pool.size(), 4);
    
    std::atomic<int> counter{0};
    std::vector<TaskCancelable::Ptr> tasks;
    
    // 提交10个任务
    for (int i = 0; i < 10; i++) {
        auto task = pool.async([&counter, i]() {
            std::this_thread::sleep_for(50ms);
            counter++;
        });
        tasks.push_back(task);
    }
    
    // 等待所有任务完成
    std::this_thread::sleep_for(600ms);
    
    // 所有任务应该已完成
    EXPECT_EQ(counter, 10);
}

// 测试任务分配到不同线程
TEST(WorkThreadPoolTest, TaskDistribution) {
    WorkThreadPool pool(4);
    
    std::mutex mutex;
    std::unordered_map<std::thread::id, int> thread_task_count;
    
    // 提交大量任务
    for (int i = 0; i < 100; i++) {
        pool.async([&mutex, &thread_task_count]() {
            auto thread_id = std::this_thread::get_id();
            std::lock_guard<std::mutex> lock(mutex);
            thread_task_count[thread_id]++;
        });
    }
    
    // 等待所有任务完成
    std::this_thread::sleep_for(500ms);
    
    // 检查任务是否分配到不同线程
    std::lock_guard<std::mutex> lock(mutex);
    EXPECT_LE(thread_task_count.size(), 4);
    EXPECT_GE(thread_task_count.size(), 1);
    
    // 输出每个线程执行的任务数量
    for (const auto& pair : thread_task_count) {
        std::cout << "Thread " << pair.first << " executed " << pair.second << " tasks" << std::endl;
    }
}

// 测试获取当前工作线程
TEST(WorkThreadPoolTest, GetCurrentWorkThread) {
    WorkThreadPool pool(2);
    
    // 在非工作线程中调用，应该返回nullptr
    auto current_thread = pool.getCurrentWorkThread();
    EXPECT_EQ(current_thread, nullptr);
    
    std::atomic<bool> got_work_thread{false};
    
    // 在工作线程中调用
    pool.async([&pool, &got_work_thread]() {
        auto work_thread = pool.getCurrentWorkThread();
        got_work_thread = (work_thread != nullptr);
        
        if (work_thread) {
            // 当前线程应该是工作线程
            EXPECT_TRUE(work_thread->isCurrentThread());
        }
    });
    
    // 等待任务执行完成
    std::this_thread::sleep_for(100ms);
    
    EXPECT_TRUE(got_work_thread);
}

// 测试获取负载最小的工作线程
TEST(WorkThreadPoolTest, GetMinLoadWorkThread) {
    WorkThreadPool pool(3);
    
    // 给第一个线程增加负载
    auto thread0 = pool.getWorkThread(0);
    auto thread1 = pool.getWorkThread(1);
    auto thread2 = pool.getWorkThread(2);
    
    std::atomic<bool> task1_running{false};
    std::atomic<bool> can_finish{false};
    
    // 提交一个长时间运行的任务给第一个线程
    thread0->async([&]() {
        task1_running = true;
        while (!can_finish) {
            std::this_thread::sleep_for(10ms);
        }
    });
    
    // 等待任务开始执行
    while (!task1_running) {
        std::this_thread::sleep_for(10ms);
    }
    
    // 获取负载最小的线程，应该不是第一个线程
    auto min_load_thread = pool.getMinLoadWorkThread();
    EXPECT_NE(min_load_thread, thread0);
    
    // 允许任务完成
    can_finish = true;
    
    // 等待任务完成
    std::this_thread::sleep_for(100ms);
}

// 测试任务取消功能
TEST(WorkThreadPoolTest, TaskCancellation) {
    WorkThreadPool pool(2);
    
    std::atomic<int> counter{0};
    
    // 提交大量任务以确保队列中有任务
    for (int i = 0; i < 20; i++) {
        pool.async([&counter]() {
            std::this_thread::sleep_for(20ms);
            counter++;
        });
    }
    
    // 提交一个任务并立即取消
    auto cancelable = pool.async([&counter]() {
        counter += 100;
    });
    
    cancelable->cancel();
    
    // 等待一段时间，确保所有任务有机会执行
    std::this_thread::sleep_for(500ms);
    
    // 计数器应该是20，如果取消的任务执行了，计数器会增加100
    EXPECT_EQ(counter, 20);
}

// 测试优先执行功能
TEST(WorkThreadPoolTest, PriorityExecution) {
    WorkThreadPool pool(1); // 单线程池，确保任务按顺序执行
    
    std::vector<int> execution_order;
    std::mutex mutex;
    
    // 提交5个普通任务
    for (int i = 0; i < 5; i++) {
        pool.async([i, &execution_order, &mutex]() {
            std::this_thread::sleep_for(50ms);
            std::lock_guard<std::mutex> lock(mutex);
            execution_order.push_back(i);
        });
    }
    
    // 提交一个优先任务
    pool.async_first([&execution_order, &mutex]() {
        std::lock_guard<std::mutex> lock(mutex);
        execution_order.push_back(999); // 特殊标记
    });
    
    // 等待所有任务完成
    std::this_thread::sleep_for(400ms);
    
    // 检查执行顺序
    std::lock_guard<std::mutex> lock(mutex);
    EXPECT_EQ(execution_order.size(), 6);
    
    // 优先任务应该在前面执行
    // 注意：由于第一个任务可能已经开始执行，优先任务可能不是第一个
    bool priority_found = false;
    for (size_t i = 0; i < execution_order.size(); i++) {
        if (execution_order[i] == 999) {
            priority_found = true;
            EXPECT_LT(i, 2); // 应该在前两个位置
            break;
        }
    }
    EXPECT_TRUE(priority_found);
}

int main(int argc, char **argv) {
    ::testing::InitGoogleTest(&argc, argv);
    return RUN_ALL_TESTS();
}